use std::{arch::asm, sync::RwLock, vec};

use lazy_static::lazy_static;

pub fn task_id() -> TaskID {
    RUNTIME.read().unwrap().task_id()
}

pub fn task_spwan(f: fn()) {   
    RUNTIME.write().unwrap().task_spawn(f);
}

pub fn task_yield() {
    RUNTIME.write().unwrap().task_yield();
}

pub fn run() {
    RUNTIME.write().unwrap().run();
}

fn task_exit() {
    RUNTIME.write().unwrap().task_exit();
}

lazy_static! {
    static ref RUNTIME: RwLock<Runtime> = {
        RwLock::new(Runtime::new())
    };
}

type TaskID = usize;

const STACK_SIZE: usize = 4096;
const TASK_SIZE: usize = 64;

struct Runtime {
    current_task: TaskID,
    tasks: Vec<Task> 
}

impl Runtime {
    fn new() -> Runtime {
        let mut main_task = Task::new(0);
        main_task.state = TaskState::Running;
        
        let mut tasks = vec![main_task];
        (1..TASK_SIZE).for_each(|i| {
            tasks.push(Task::new(i));
        });
        
        Runtime {
            current_task: 0,
            tasks,
        }
    }

    fn task_id(&self) -> TaskID {
        self.current_task
    }

    fn task_spawn(&mut self, f: fn()) {
        // Find a Available task
        let task = self.tasks.iter_mut()
            .find(|task| task.state == TaskState::Available)
            .expect("No more task!");

        // Init task
        task.state = TaskState::Ready;
        task.context = TaskContext::default();
        // All task return to 'coroutine_return'
        task.context.x1 = task_exit as u64;
        // Init stack point
        task.context.x2 = task.stack_point();
        // Task function entry
        task.context.nx1 = f as u64;
    }

    fn task_yield(&mut self) -> bool {
        // Find a Ready task
        let mut next_task = self.current_task;
        while self.tasks[next_task].state != TaskState::Ready {
            next_task += 1;
            if next_task == self.tasks.len() {
                next_task = 0;
            }
            if next_task == self.current_task {
                return false;
            }
        }

        self.tasks[self.current_task].state = TaskState::Ready;
        self.tasks[next_task].state = TaskState::Running;

        let last_task = self.current_task;
        self.current_task = next_task;

        unsafe {
            switch(
                &mut self.tasks[last_task].context, 
                &self.tasks[next_task].context
            );
        }

        true
    }

    fn task_exit(&mut self) {
        if self.current_task != 0 {
            self.tasks[self.current_task].state = TaskState::Available;
            self.task_yield();
        }
    }

    fn run(&mut self) {
        while self.task_yield() {}
        std::process::exit(0);
    }
}

#[inline(never)]
unsafe fn switch(_last_task: *mut TaskContext, _next_task: *const TaskContext) {
    // a0: last, a1: next
    asm!("
        sd x1, 0x00(a0)
        sd x2, 0x08(a0)
        sd x8, 0x10(a0)
        sd x9, 0x18(a0)
        sd x18, 0x20(a0)
        sd x19, 0x28(a0)
        sd x20, 0x30(a0)
        sd x21, 0x38(a0)
        sd x22, 0x40(a0)
        sd x23, 0x48(a0)
        sd x24, 0x50(a0)
        sd x25, 0x58(a0)
        sd x26, 0x60(a0)
        sd x27, 0x68(a0)
        sd x1, 0x70(a0)

        ld x1, 0x00(a1)
        ld x2, 0x08(a1)
        ld x8, 0x10(a1)
        ld x9, 0x18(a1)
        ld x18, 0x20(a1)
        ld x19, 0x28(a1)
        ld x20, 0x30(a1)
        ld x21, 0x38(a1)
        ld x22, 0x40(a1)
        ld x23, 0x48(a1)
        ld x24, 0x50(a1)
        ld x25, 0x58(a1)
        ld x26, 0x60(a1)
        ld x27, 0x68(a1)
        ld t0, 0x70(a1)

        jr t0
    ")
}

#[allow(unused)]
struct Task {
    id: TaskID,
    state: TaskState,
    context: TaskContext,
    stack: Vec<u8>,
}

impl Task {
    fn new(id: TaskID) -> Self {
        Self {
            id,
            state: TaskState::Available,
            context: TaskContext::default(),
            stack: vec![0u8; STACK_SIZE],
        }
    }

    fn stack_point(&self) -> u64 {
        self.stack.as_ptr() as *const _ as u64
    }
}

#[derive(Debug, PartialEq, Eq)]
enum TaskState {
    Available,
    Running,
    Ready,
}

#[derive(Debug, Default)]
#[repr(C)]
struct TaskContext {
    // 15 u64
    x1: u64,  //ra: return addres
    x2: u64,  //sp
    x8: u64,  //s0,fp
    x9: u64,  //s1
    x18: u64, //x18-27: s2-11
    x19: u64,
    x20: u64,
    x21: u64,
    x22: u64,
    x23: u64,
    x24: u64,
    x25: u64,
    x26: u64,
    x27: u64,
    nx1: u64, //new return addres
}

